import os
import platform
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from datetime import datetime
from math import ceil

import fitz
import pymupdf


# from comm.common import log


# 根据任务类型和硬件架构设置进程池进程数量
def get_optimal_worker_count(task_type):
    cpu_count = os.cpu_count()
    machine = platform.machine()

    # print(f"machine={machine},cpu_count={cpu_count}")

    if machine == "arm64":  # 对于 Apple Silicon
        if task_type == "cpu_intensive":
            return max(1, cpu_count // 2)  # 使用高性能核心的数量
        elif task_type == "io_intensive":
            return max(1, cpu_count - 2)  # 尝试使用更多的进程
        else:
            return max(1, cpu_count // 2)  # 默认情况下使用高性能核心的数量
    else:  # 对于 Intel 或其他架构
        if task_type == "cpu_intensive":
            return cpu_count  # 对于 CPU 密集型任务，使用 CPU 核心数量
        elif task_type == "io_intensive":
            return cpu_count * 2  # 对于 I/O 密集型任务，试试 CPU 核心数量的 2 倍
        else:
            return cpu_count  # 默认情况下，使用 CPU 核心数量


def convert_page_to_image(page_number, pdf_path, output_folder, resolution, color_mode, type, alpha, image_name_prefix):
    pdf_document = fitz.open(pdf_path)
    page = pdf_document[page_number]

    # 获取页面的原始大小（以点为单位）
    zoom_x = resolution / page.rect.width
    zoom_y = resolution / page.rect.height

    # 使用最小的缩放比例以确保整个页面都适应目标分辨率
    zoom = min(zoom_x, zoom_y)

    # 创建一个矩阵来应用缩放和抗锯齿
    mat = fitz.Matrix(zoom, zoom)

    # 将 PDF 页面转换为图片，指定分辨率、颜色模式和抗锯齿
    pix = page.get_pixmap(matrix=mat, alpha=alpha, colorspace=color_mode)

    # 定义输出图片的路径
    output_path = f'{output_folder}/{image_name_prefix}{page.number + 1}.{type}'

    # 保存图片
    pix.save(output_path)


def pdf_to_images(pdf_path, output_folder, resolution=2048, color_mode='rgb', img_type='png', is_alpha=True,
                  image_name_prefix=None, task_type="cpu_intensive"):
    ret_is_ok = False
    try:
        # 检查图片格式是否正确
        if img_type not in ['png', 'pnm', 'pgm', 'ppm', 'pbm', 'pam', 'psd', 'ps', 'jpg', 'jpeg']:
            print(f"[Error] Image format '{img_type}' cannot be parsed")
            return ret_is_ok, "[Error] format error"

        # 检查清晰度是否为0
        if resolution == 0:
            print("[Error] The resolution parameter cannot be 0")
            return ret_is_ok, "[Error] resolution equals 0"

        # 检查图片格式是否支持 alpha 通道
        if img_type in ['pbm', 'pgm', 'ppm', 'pnm'] and is_alpha:
            print(f"[Error] The '{img_type}' format cannot have alpha")
            return ret_is_ok, "[Error] cannot have alpha"

        # 打开 PDF 文件
        pdf_document = fitz.open(pdf_path)

        # 如果没有提供 image_name_prefix 参数，使用当前时间作为默认值
        if image_name_prefix is None:
            image_name_prefix = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")

        # 检查输出文件夹是否存在，如果不存在，则创建一个
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

        # 检查PDF页数
        pdf_page_count = len(pdf_document)
        works = get_optimal_worker_count(task_type)
        tasks = ceil(pdf_page_count / works)

        # 创建一个进程池
        with ProcessPoolExecutor(max_workers=works) as executor:
            # 遍历 PDF 中的每一页
            futures = []
            for i in range(tasks):
                start = i * works
                end = (i + 1) * works
                if end > pdf_page_count:
                    end = pdf_page_count
                for page_number in range(start, end):
                    futures.append(executor.submit(convert_page_to_image, page_number, pdf_path, output_folder,
                                                   resolution, color_mode, img_type, is_alpha, image_name_prefix))
                if end == pdf_page_count:
                    break

            for future in as_completed(futures):
                future.result()

        # 关闭 PDF 文件
        pdf_document.close()

        ret_is_ok = True
        print("Successfully converted!")
        return ret_is_ok, "successfully converted"
    except pymupdf.FileNotFoundError:
        print("[Error] file does not exist")
        return ret_is_ok, "[Error] file does not exist"
    except Exception as e:
        print(f"[Error] {e}")
        return ret_is_ok, f"[Error] {e}"


if __name__ == '__main__':
    pdf_path = 'F:/法智通/第二期第二批/hungeerjincheng_second_02_240911.7z/资产材料20240911/5.长银消金/长银消金（中智信）-齐/长银消金-代偿证明（中智信与象源）/长银消金-债转协议（中智信到象源）.pdf'
    output_folder = 'images/test_file/test_file_gray/'

    start_time = time.time()
    pdf_to_images(pdf_path=pdf_path, output_folder=output_folder, resolution=2048, color_mode='rgb', img_type='png',
                  is_alpha=True, image_name_prefix='test_file_gray')
    print(f"运行时间: {time.time() - start_time}")
